import os
from enum import Enum

from langchain.output_parsers import DatetimeOutputParser, EnumOutputParser
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

#llm = Ollama(model="llama3", base_url="http://localhost:11434")
llm = ChatOpenAI(model="gpt-3.5-turbo", api_key=os.environ["OPENAI_API_KEY_ZHIHU"],base_url=os.environ["OPENAI_API_BASE_ZHIHU"])
# 定义提示模板
class Color(Enum):
    RED = "Red"
    GREEN = "Green"
    BLUE = "Blue"
    BLACK = "Black"
    BROWN = "Brown"
    WHITE = "White"

parser  = EnumOutputParser(enum=Color)

prompt = PromptTemplate(template="What color does this animal look like?\nAnimal: {animal}\n{format_instructions}", input_variables=["animal"],
                        partial_variables = {"format_instructions":parser.get_format_instructions()})

print(prompt)
chain = prompt | llm | parser
response = chain.invoke({"animal":"北极熊"})
print(response)



